"""
  这份文件用来读取数据
"""

# 其他常量
FOLLOW_TARGET_FRAMES = 200 # 跟随状态的目标帧数（要收集的数据数量）
LANECHANGE_TARGET_FRAMES = 80 # 车道变换状态的目标帧数（要收集的数据数量）
FREE_DRIVING_FRAMES = 150 # 自由驾驶状态的目标帧数（要收集的数据数量）
FREE_DRIVING_PREDICT_FRAMES = 10 # 自由驾驶状态的预测数据帧（要收集的数据数量）

# 读取数据，并基于CarId分组（TODO 废弃）
def read_data_and_group_by_carId(fileName, self_carId_idx=0, is_read_follow_data=False, is_read_lanechange_data=False, front_carId_idx=1, after_carId_idx=3):
    # 收集相同ID数车辆数据
    cars_with_commonId = []
    # 结果数组
    res = []

    with open(fileName, "r") as file_reader:
        # 先读一行
        cars_with_commonId = [[float(i) for i in file_reader.readline().split()]]

        line_str = ""

        # 持续性读取一行数据
        while line_str := file_reader.readline():
            # 字符串 to 浮点型数组
            line_float_arr = [float(i) for i in line_str.split()]

            if(cars_with_commonId[0][self_carId_idx] == line_float_arr[self_carId_idx]):
                # 自车同ID
                if is_read_follow_data:
                    # !读取的是 提取的跟驰数据
                    if cars_with_commonId[0][front_carId_idx] == line_float_arr[front_carId_idx] and len(cars_with_commonId) < FOLLOW_TARGET_FRAMES:
                        # 前车同ID，并且没有超出帧数限制
                        cars_with_commonId.append(line_float_arr)
                        continue
                    
                    # 收集到一组完整数据
                    res.append(cars_with_commonId)
                    # 新ID
                    cars_with_commonId = [line_float_arr]
                    continue

                if is_read_lanechange_data:
                    # !读取的是 提取的换道数据
                    if cars_with_commonId[0][front_carId_idx] == line_float_arr[front_carId_idx] and cars_with_commonId[0][after_carId_idx] == line_float_arr[after_carId_idx] and len(cars_with_commonId) < LANECHANGE_TARGET_FRAMES:
                        # 前后车都同ID，并且没有超出帧数限制
                        cars_with_commonId.append(line_float_arr)
                        continue
                
                # !读取的是 原始数据 或者 过滤数据
                cars_with_commonId.append(line_float_arr)
                continue
 
            # 收集到一组完整数据
            res.append(cars_with_commonId)
            # 新ID
            cars_with_commonId = [line_float_arr]
    
    #! 防止漏掉最后一组数据  
    if len(cars_with_commonId) > 0:
        res.append(cars_with_commonId)
    
    return res

# 读取原始数据、过滤的数据
def read_original_data(fileName, self_carId_idx=0):
    # 收集相同ID数车辆数据
    cars_with_commonId = []
    # 结果数组
    res = []

    with open(fileName, "r") as file_reader:
        # 先读一行
        cars_with_commonId = [[float(i) for i in file_reader.readline().split()]]

        line_str = ""
        # 持续性读取一行数据
        while line_str := file_reader.readline():
            # 字符串 to 浮点型数组
            line_float_arr = [float(i) for i in line_str.split()]

            if(cars_with_commonId[0][self_carId_idx] == line_float_arr[self_carId_idx]):
                cars_with_commonId.append(line_float_arr)
            else:
                # 收集到一组完整数据
                res.append(cars_with_commonId)
                # 新ID
                cars_with_commonId = [line_float_arr]
    
    #! 防止漏掉最后一组数据  
    if len(cars_with_commonId) > 0:
        res.append(cars_with_commonId)
    
    return res

# 读取提取的跟驰数据
def read_extract_follow_data(fileName, front_carId_idx=1, self_carId_idx=2):
    # 收集相同ID数车辆数据
    cars_with_commonId = []
    # 结果数组
    res = []

    with open(fileName, "r") as file_reader:
        # 先读一行
        cars_with_commonId = [[float(i) for i in file_reader.readline().split()]]

        line_str = ""
        # 持续性读取一行数据
        while line_str := file_reader.readline():
            # 字符串 to 浮点型数组
            line_float_arr = [float(i) for i in line_str.split()]

            if(cars_with_commonId[0][self_carId_idx] == line_float_arr[self_carId_idx]):
                # 自车同ID
                if cars_with_commonId[0][front_carId_idx] == line_float_arr[front_carId_idx] and len(cars_with_commonId) < FOLLOW_TARGET_FRAMES:
                    # 前车同ID，并且没有超出帧数限制
                    cars_with_commonId.append(line_float_arr)
                else:
                    # 收集到一组完整数据
                    res.append(cars_with_commonId)
                    # 新ID
                    cars_with_commonId = [line_float_arr]
            else:
                # 收集到一组完整数据
                res.append(cars_with_commonId)
                # 新ID
                cars_with_commonId = [line_float_arr]
    
    #! 防止漏掉最后一组数据  
    if len(cars_with_commonId) > 0:
        res.append(cars_with_commonId)

    return res

# 读取提取的换道数据
def read_extract_lanechange_data(fileName, front_carId_idx=1, self_carId_idx=2, after_carId_idx=3):
    # 收集相同ID数车辆数据
    cars_with_commonId = []
    # 结果数组
    res = []

    with open(fileName, "r") as file_reader:
        # 先读一行
        cars_with_commonId = [[float(i) for i in file_reader.readline().split()]]

        line_str = ""
        # 持续性读取一行数据
        while line_str := file_reader.readline():
            # 字符串 to 浮点型数组
            line_float_arr = [float(i) for i in line_str.split()]

            if(cars_with_commonId[0][self_carId_idx] == line_float_arr[self_carId_idx]):
                # 自车同ID
                if cars_with_commonId[0][front_carId_idx] == line_float_arr[front_carId_idx] and cars_with_commonId[0][after_carId_idx] == line_float_arr[after_carId_idx] and len(cars_with_commonId) < LANECHANGE_TARGET_FRAMES:
                    # 前后车都同ID，并且没有超出帧数限制
                    cars_with_commonId.append(line_float_arr)
                else:
                    # 收集到一组完整数据
                    res.append(cars_with_commonId)
                    # 新ID
                    cars_with_commonId = [line_float_arr]
            else:
                # 收集到一组完整数据
                res.append(cars_with_commonId)
                # 新ID
                cars_with_commonId = [line_float_arr]
    
    #! 防止漏掉最后一组数据  
    if len(cars_with_commonId) > 0:
        res.append(cars_with_commonId)
    
    return res

# 读取提取的自由行驶数据
def read_extract_free_drive_data(fileName):
  # 结果数组
  res = []
  with open(fileName, "r") as file_reader:
    group_data = []
    line_str = ""
    cur_frames = 1

    # 持续性读取一行
    while line_str := file_reader.readline():
      # 字符串 to 浮点型数组
      line_float_arr = [float(i) for i in line_str.split()]

      if(cur_frames <= FREE_DRIVING_FRAMES):
        # 帧数不够
        group_data.append(line_float_arr)
        cur_frames += 1
      else:
        # 达到目标帧数
        res.append(group_data)
        group_data = [line_float_arr]
        cur_frames = 2

      #! 防止漏掉最后一组数据  
    if len(group_data) > 0:
      res.append(group_data)
  
  return res

# 读取lstm模型验证数据
def read_lstm_verify_data(fileName, targetFrames=10):
    data = []
    with open(fileName, "r") as file_reader:
      cur_frames = 0
      group = []
      while line_str := file_reader.readline():
        group.append([float(i) for i in line_str.split()])
        cur_frames += 1

        # 帧数检查
        if cur_frames >= targetFrames:
           data.append(group)
           group = []
           cur_frames = 0

    if len(group) > 0:
      data.append(group)
    
    return data
